# -*- coding: utf-8 -*-
"""
SHP文件处理器
将SHP文件转换为标准化的PostGIS导入参数格式
"""

import geopandas as gpd
import pandas as pd
from typing import Dict, Any, Optional
import os
import json
from loguru import logger
from shapely.geometry import MultiPolygon, Point


class SHPProcessor:
    """SHP文件处理器，转换为PostGIS导入格式"""
    
    def __init__(self, file_path: str):
        self.file_path = file_path
        self.gdf = None
        
    def load_data(self, encoding: str = 'utf-8') -> Dict[str, Any]:
        """
        加载SHP文件并转换为标准格式
        
        Returns:
            Dict包含: {
                'data': 标准化后的GeoDataFrame,
                'geometry_type': 几何类型,
                'crs': 坐标系,
                'columns': 字段列表,
                'record_count': 记录数量,
                'feature_count': 原始特征数量
            }
        """
        try:
            # 读取SHP文件
            self.gdf = gpd.read_file(self.file_path, encoding=encoding)
            
            # 标准化数据结构
            standardized_data = self._standardize_data()
            
            result = {
                'data': standardized_data,
                'geometry_type': self._get_geometry_type(),
                'crs': str(self.gdf.crs) if self.gdf.crs else 'EPSG:4326',
                'columns': list(standardized_data.columns),
                'record_count': len(standardized_data),
                'feature_count': len(self.gdf),
                'source_format': 'SHP'
            }
            
            logger.info(f"SHP处理完成: {os.path.basename(self.file_path)}")
            logger.info(f"特征数: {result['feature_count']}, 几何类型: {result['geometry_type']}")
            
            return result
            
        except Exception as e:
            logger.error(f"SHP文件处理失败: {e}")
            raise
    
    def _standardize_data(self) -> gpd.GeoDataFrame:
        """标准化数据结构"""
        if self.gdf is None:
            raise ValueError("请先加载数据")
            
        # 创建标准化列
        standardized = self.gdf.copy()
        
        # SHP文件通常不包含嵌套的properties字段，这是GeoJSON特有的
        
        # 确保有标准字段
        if 'code' not in standardized.columns:
            # 优先尝试stcd字段作为code，然后尝试其他常见水文字段
            code_candidates = ['stcd', 'STCD', 'stcd_field', 'gd', 'GD', 'id', 'ID', 'code', 'Code', 'CODE', 'station_id', 'StationID']
            for candidate in code_candidates:
                if candidate in standardized.columns:
                    standardized['code'] = standardized[candidate].astype(str)
                    logger.info(f"使用字段 '{candidate}' 作为code")
                    break
            else:
                standardized['code'] = [f"SHP_{i:06d}" for i in range(len(standardized))]
                logger.warning("未找到合适的code字段，使用自动生成")

        if 'name' not in standardized.columns:
            # 尝试从更多可能的名称字段获取
            name_candidates = ['name', 'Name', 'NAME', 'stnm', 'STNM', 'title', 'label', 'station_name', 'reservoir_name', 'lake_name']
            for candidate in name_candidates:
                if candidate in standardized.columns:
                    standardized['name'] = standardized[candidate].astype(str)
                    logger.info(f"使用字段 '{candidate}' 作为name")
                    break
            else:
                # 如果没有合适的名称字段，尝试使用code作为name
                if 'code' in standardized.columns:
                    standardized['name'] = standardized['code']
                    logger.info("使用code作为name")
                else:
                    standardized['name'] = [f"未命名_{i}" for i in range(len(standardized))]
                
        if 'type' not in standardized.columns:
            # 尝试从类型相关字段获取
            type_candidates = ['type', 'Type', 'TYPE', 'category', 'Category', 'kind', 'Kind', 'reservoir_type', 'station_type']
            for candidate in type_candidates:
                if candidate in standardized.columns:
                    standardized['type'] = standardized[candidate].astype(str)
                    logger.info(f"使用字段 '{candidate}' 作为type")
                    break
            else:
                # 根据几何类型推断合理的默认值
                first_geom_type = standardized.geometry.iloc[0].geom_type if len(standardized) > 0 else 'Unknown'
                if first_geom_type in ['Point', 'MultiPoint']:
                    standardized['type'] = '水文站点'
                elif first_geom_type in ['Polygon', 'MultiPolygon']:
                    standardized['type'] = '水库'
                else:
                    standardized['type'] = '要素'
            
        # 多余字段直接丢弃，extra_data保持为空
        standardized['extra_data'] = ''
        
        # 转换几何类型以匹配数据库要求
        standardized = self._convert_geometry_types(standardized)
        
        # 转换坐标系为EPSG:4326
        if standardized.crs is None:
            standardized = standardized.set_crs('EPSG:4326')
        elif str(standardized.crs) != 'EPSG:4326':
            standardized = standardized.to_crs('EPSG:4326')
            
        # 只保留标准字段
        result = standardized[['code', 'name', 'type', 'extra_data', 'geometry']].copy()
        
        return result
    
    def _convert_geometry_types(self, gdf: gpd.GeoDataFrame) -> gpd.GeoDataFrame:
        """转换几何类型以匹配数据库要求"""
        if gdf is None or len(gdf) == 0:
            return gdf
            
        # 创建副本避免修改原始数据
        result = gdf.copy()
        
        # 获取原始几何类型
        first_geom_type = result.geometry.iloc[0].geom_type
        
        # 根据类型进行转换
        if first_geom_type == 'Polygon':
            # 将Polygon转换为MultiPolygon
            result.geometry = result.geometry.apply(
                lambda geom: geom if geom is None else 
                geom if geom.geom_type == 'MultiPolygon' else
                geom if geom.geom_type != 'Polygon' else
                MultiPolygon([geom])
            )
        elif first_geom_type == 'MultiPoint':
            # 将MultiPoint转换为Point（取第一个点）
            result.geometry = result.geometry.apply(
                lambda geom: geom if geom is None else
                geom if geom.geom_type == 'Point' else
                geom if geom.geom_type != 'MultiPoint' else
                list(geom.geoms)[0] if len(list(geom.geoms)) > 0 else None
            )
            
        return result
    
    def _get_geometry_type(self) -> str:
        """获取几何类型"""
        if self.gdf is None or len(self.gdf) == 0:
            return 'Unknown'
            
        geom_type = self.gdf.geometry.iloc[0].geom_type
        
        # 标准化几何类型
        type_mapping = {
            'Polygon': 'MULTIPOLYGON',
            'MultiPolygon': 'MULTIPOLYGON',
            'Point': 'POINT',
            'MultiPoint': 'POINT'
        }
        
        return type_mapping.get(geom_type, geom_type.upper())
    
    def validate_data(self) -> Dict[str, Any]:
        """验证数据完整性"""
        if self.gdf is None:
            return {'valid': False, 'errors': ['数据未加载']}
            
        errors = []
        warnings = []
        
        # 检查几何有效性
        invalid_geoms = ~self.gdf.geometry.is_valid
        if invalid_geoms.any():
            errors.append(f"发现{invalid_geoms.sum()}个无效几何图形")
            
        # 检查空几何
        empty_geoms = self.gdf.geometry.is_empty
        if empty_geoms.any():
            errors.append(f"发现{empty_geoms.sum()}个空几何图形")
            
        # 检查重复code
        if 'code' in self.gdf.columns:
            duplicate_codes = self.gdf['code'].duplicated().sum()
            if duplicate_codes > 0:
                errors.append(f"发现{duplicate_codes}个重复code")
                
        # 检查坐标系
        if self.gdf.crs is None:
            warnings.append("未定义坐标系，将使用EPSG:4326")
            
        return {
            'valid': len(errors) == 0,
            'errors': errors,
            'warnings': warnings
        }

    def get_metadata(self) -> Dict[str, Any]:
        """获取文件元数据"""
        if self.gdf is None:
            return {}
            
        return {
            'bbox': self.gdf.total_bounds.tolist(),
            'crs': str(self.gdf.crs),
            'geometry_types': self.gdf.geometry.geom_type.value_counts().to_dict(),
            'field_types': self.gdf.dtypes.to_dict()
        }


# 测试示例
if __name__ == "__main__":
    # 测试SHP文件处理
    test_file = "水库合并shp/merged_output.shp"  # 示例路径
    
    if os.path.exists(test_file):
        processor = SHPProcessor(test_file)
        
        # 加载并处理数据
        result = processor.load_data()
        
        # 验证数据
        validation = processor.validate_data()
        
        print("=== SHP处理结果 ===")
        print(f"文件: {test_file}")
        print(f"记录数: {result['record_count']}")
        print(f"几何类型: {result['geometry_type']}")
        print(f"坐标系: {result['crs']}")
        print(f"字段: {result['columns']}")
        print(f"数据验证: {'通过' if validation['valid'] else '失败'}")
        
        if not validation['valid']:
            print("错误:", validation['errors'])
            
        # 显示前几条数据
        print("\n=== 前3条数据示例 ===")
        print(result['data'].head(3)[['code', 'name', 'type']])
    else:
        print(f"测试文件不存在: {test_file}")
        print("请使用实际存在的SHP文件路径进行测试")